#include <ipa_building_navigation/ipa_room_planning.h>


IpaRoomPlanning::IpaRoomPlanning(rclcpp::Node::SharedPtr &nh,std::shared_ptr<tf2_ros::Buffer> tf_)
{
    node_handle_=nh;
    tf_buffer=tf_;
    all_path_pub = node_handle_->create_publisher<nav_msgs::msg::Path>("all_coverage_path", 2);
    room_segmentation_client= std::make_shared<SimpleActionClientD<ClientSegmentation>>(node_handle_,std::string("room_segmentation_server"));
    room_sequence_client=     std::make_shared<SimpleActionClientD<ClientSequence>>(node_handle_,std::string("room_sequence_planning_server"));
    room_exploration_client=  std::make_shared<SimpleActionClientD<ClientExploration>>(node_handle_,std::string("room_exploration_server"));
    grid_client = node_handle_->create_client<nav2_msgs::srv::GetMapNav2>("/map_server/map");
    free_thresh= 0.25;
    occupied_thresh= 0.65;
}


bool IpaRoomPlanning::prepareData()
{
    std::vector<std::vector<bool> > grid;
    std::vector<double> map_origin(3,0);
    std::vector<double> robot_pos(3,0);
    std::vector<geometry_msgs::msg::Point32> fov_points(4);
    geometry_msgs::msg::Point32 fov_origin;
    fov_points[0].x = -ipa_planning_param_.robot_radius;		
	fov_points[0].y = ipa_planning_param_.robot_radius;
	fov_points[1].x = -ipa_planning_param_.robot_radius;
	fov_points[1].y = -ipa_planning_param_.robot_radius;
	fov_points[2].x = ipa_planning_param_.robot_radius;
	fov_points[2].y = -ipa_planning_param_.robot_radius;
	fov_points[3].x = ipa_planning_param_.robot_radius;
	fov_points[3].y = ipa_planning_param_.robot_radius;
	fov_origin.x = 0.;
	fov_origin.y = 0.;
    ipa_planning_param_.fov_points= fov_points;
    ipa_planning_param_.fov_origin=fov_origin;
    ipa_planning_param_.start_pos=robot_pos;

	auto grid_req_srv =  std::make_shared<nav2_msgs::srv::GetMapNav2::Request>();
	std::cout<<"Requesting grid!"<< std::endl;
    if(!grid_client->wait_for_service(std::chrono::seconds(5)))
	{
        std::cout<<"Could not retrieve grid from map_server\n"<< std::endl;
        try {
            std::cout << "data_source: default path image !" <<std::endl;
            std::string image_file = 
            "/home/ora/ros2_ws/src/Sweeping_robot/ipa_coverage_planning/ipa_room_segmentation/common/files/test_maps/map2.pgm";
            ipa_planning_param_.image_file=image_file;
            fromPathGetMat(ipa_planning_param_.image_file);
        }
        catch(...)
        {
            std::cout << "no data source !" <<std::endl;
            return false;
        }
	}
    else
	{
        auto result_future= grid_client->async_send_request(grid_req_srv);
        if (rclcpp::spin_until_future_complete(node_handle_, result_future) ==
            rclcpp::FutureReturnCode::SUCCESS)
        {
            std::cout<<"Requesting grid  success!"<< std::endl;
            auto srv_goal = result_future.get();
            if(srv_goal->data_source== srv_goal->IMAGE_PATH)
            {
                std::cout << "data_source: IMAGE_PATH" <<std::endl;
                ipa_planning_param_.image_file=srv_goal->map_path;
                fromPathGetMat(srv_goal->map_path);
            }
            else if(srv_goal->data_source== srv_goal->IMAGE_DATA)
            {
                std::cout << "data_source: IMAGE_DATA" <<std::endl;
                cv_bridge::CvImagePtr cv_ptr_obj;
                input_image_msg=srv_goal->input_map;
                cv_ptr_obj = cv_bridge::toCvCopy(input_image_msg, sensor_msgs::image_encodings::MONO8);
                cv_ptr_obj->image.copyTo(origin_map_mat);
                Binarization(origin_map_mat,250);
            }
            else if(srv_goal->data_source== srv_goal->MAP_DATA)
            {
                std::cout << "data_source: MAP_DATA" <<std::endl;
                map_origin[0]= srv_goal->map_origin.position.x;
                map_origin[1]= srv_goal->map_origin.position.y;
                ipa_planning_param_.map_resolution=srv_goal->map_resolution;
                mapToMat(srv_goal->map, origin_map_mat);
                cv::flip(origin_map_mat, origin_map_mat, 0);
                Binarization(origin_map_mat,250);
                cv_bridge::CvImage cv_image;
                cv_image.encoding = "mono8";
                cv_image.image = origin_map_mat;
                cv_image.toImageMsg(input_image_msg);
            }
            else
                return false;
        } else {
            RCLCPP_ERROR(rclcpp::get_logger("rclcpp"), "Failed to call service add_two_ints");
            return false;
        }
    }
    if(getRobotPose(tf_buffer,robot_pos))
    {
        ipa_planning_param_.robot_pos=robot_pos;
        ipa_planning_param_.start_pos=robot_pos;
    }
	
    ipa_planning_param_.map_origin=map_origin;
    return true;
}

void IpaRoomPlanning::fromPathGetMat(std::string image_path)
{
	origin_map_mat = cv::imread(image_path, 0);
    Binarization(origin_map_mat,250);
	std::cout << "map-size: " << origin_map_mat.rows << "x" << origin_map_mat.cols << std::endl;
	cv_bridge::CvImage cv_image;
	cv_image.encoding = "mono8";
	cv_image.image = origin_map_mat;
	cv_image.toImageMsg(input_image_msg);
}

void IpaRoomPlanning::matToMap(nav_msgs::msg::OccupancyGrid &map, const cv::Mat &mat)
{
	map.info.width  = mat.cols;
	map.info.height = mat.rows;
	map.data.resize(mat.cols*mat.rows);

	for(int x=0; x<mat.cols; x++)
		for(int y=0; y<mat.rows; y++)
			map.data[y*mat.cols+x] = mat.at<int8_t>(y,x)?0:100;
}

void IpaRoomPlanning::mapToMat(const nav_msgs::msg::OccupancyGrid &map, cv::Mat &mat)
{
	mat = cv::Mat(map.info.height, map.info.width, CV_8U);
	int free_thresh_int = std::rint(free_thresh * 100.0);
    int occupied_thresh_int = std::rint(occupied_thresh * 100.0);
	double map_cell;
	int8_t value;
	for(int x=0; x<mat.cols; x++)
		for(int y=0; y<mat.rows; y++)
		{
			map_cell=map.data[y*mat.cols+x];
			if (map_cell < 0 || 100 < map_cell) {
			  value = 0; //205
            }
			else if (map_cell <= free_thresh_int) {
			  value = 255;
            } 
			else if (occupied_thresh_int <= map_cell) {
			  value = 0;
            } 
			else {
			  value = 128;  // 
            }
			mat.at<int8_t>(y,x) =value;
		}
}

void IpaRoomPlanning::Binarization(cv::Mat &map_copy,int threshold,int min_value,int max_value)
{
    for (int y = 0; y < map_copy.rows; y++)
    {
        for (int x = 0; x < map_copy.cols; x++)
        {
            if (map_copy.at<unsigned char>(y, x) >= threshold)
            {
                map_copy.at<unsigned char>(y, x) = max_value;
            }
            else
            {
                map_copy.at<unsigned char>(y, x) = min_value;
            }
        }
    }
}
void IpaRoomPlanning::Publish( )
{
    nav_msgs::msg::Path path_msg;
    path_msg.header.frame_id = "map";
    path_msg.header.stamp = rclcpp::Time();
    path_msg.poses.clear();
    for(int i=0; i< all_room_paths.size();i++)
    {
        path_msg.poses.insert(path_msg.poses.end(),all_room_paths[i].begin(),all_room_paths[i].end());
    }
    all_path_pub->publish(path_msg);
    all_room_paths.clear();
}


bool IpaRoomPlanning::getRoomExploration(IpaRoomPlanningParam & exploration_param,cv::Mat & exploration,cv::Mat & display_mat)
{
    if(exploration_param.do_flip)
    {
         cv::flip(exploration, exploration, 0);
    }
    Binarization(exploration,1);

    sensor_msgs::msg::Image labeling;
    cv_bridge::CvImage cv_image;
    cv_image.encoding = "mono8";
    cv_image.image = exploration;
    cv_image.toImageMsg(labeling);
    ipa_building_msgs::action::RoomExploration_Goal goal;
    goal.input_map = labeling;
    goal.map_resolution = exploration_param.map_resolution;
    goal.map_origin.position.x = exploration_param.map_origin[0];
    goal.map_origin.position.y = exploration_param.map_origin[1];
    goal.map_origin.position.z = exploration_param.map_origin[2];
    goal.robot_radius = exploration_param.robot_radius; // turtlebot, used for sim 0.177, 0.4
    goal.coverage_radius = exploration_param.coverage_radius;
    goal.field_of_view = exploration_param.fov_points;
    goal.field_of_view_origin = exploration_param.fov_origin;
    goal.starting_position.x = exploration_param.start_pos[0];
    goal.starting_position.y = exploration_param.start_pos[1];
    goal.starting_position.theta = exploration_param.start_pos[2];
    goal.planning_mode = exploration_param.planning_mode;
    while(rclcpp::ok() && !room_exploration_client->wait_server(5)) 
    {
        std::cout<<"room_exploration_client Action server not available after waiting, wait 5s..."<< std::endl;
    }
    printf("room_exploration_client, sending goal.\n");
    result_exp= nullptr;
    result_exp= room_exploration_client->sendGoal(goal);
    if(result_exp== nullptr)
    {
        printf("Timeout on room segmentation.");
        return false;
    }
    all_room_paths.push_back(result_exp->coverage_path_pose_stamped);
    printf("room exploration successfully!");

    if(exploration_param.display)
    {
        if(display_mat.cols!= colour_segmented_map.cols || display_mat.rows!= colour_segmented_map.rows)
        {
            display_mat = colour_segmented_map.clone();
            if(exploration_param.do_flip)
                cv::flip(display_mat, display_mat, 0);
        }
        
        const double inverse_map_resolution = 1./goal.map_resolution;
        for (size_t point=0; point<result_exp->coverage_path.size(); ++point)
        {
            const cv::Point point1((result_exp->coverage_path[point].x-exploration_param.map_origin[0])*inverse_map_resolution,
                    (result_exp->coverage_path[point].y-exploration_param.map_origin[1])*inverse_map_resolution);
            cv::circle(display_mat, point1, 2, cv::Scalar(128), cv::FILLED);
            if (point > 0)
            {
                const cv::Point point2((result_exp->coverage_path[point-1].x-exploration_param.map_origin[0])*inverse_map_resolution,
                         (result_exp->coverage_path[point-1].y-exploration_param.map_origin[1])*inverse_map_resolution);
                cv::line(display_mat, point1, point2, cv::Scalar(128), 1);
            }
            //std::cout << "coverage_path[" << point << "]: x=" << result_exp->coverage_path[point].x << ", y=" << result_exp->coverage_path[point].y << ", theta=" << result_exp->coverage_path[point].theta << std::endl;
        }
        //cv::imshow("exploration client", display_mat);
        //cv::waitKey(0);
    }
    return true;
}

bool IpaRoomPlanning::getRoomSequence(IpaRoomPlanningParam & sequence_param)
{
    std::cout<< "run getRoomSequence" << std::endl;
    while(rclcpp::ok() && !room_sequence_client->wait_server(5)) 
    {
        std::cout<<"room_sequence_client Action server not available after waiting, wait 5s..."<< std::endl;
    }
    printf("room_sequence_client, sending goal.\n");
    ipa_building_msgs::action::FindRoomSequenceWithCheckpoints_Goal goal_seq;
    goal_seq.input_map = input_image_msg;
    goal_seq.map_resolution = sequence_param.map_resolution;
    goal_seq.map_origin.position.x = sequence_param.map_origin[0];
    goal_seq.map_origin.position.y = sequence_param.map_origin[1];
    goal_seq.map_origin.position.z = sequence_param.map_origin[2];
    goal_seq.room_information_in_pixel = result_seg->room_information_in_pixel;
    goal_seq.robot_radius = sequence_param.robot_radius;
    bool flag=false;
    cv::Mat map_eroded;
    cv::erode(origin_map_mat, map_eroded, cv::Mat(), cv::Point(-1,-1), goal_seq.robot_radius/goal_seq.map_resolution+2);
    cv::Mat distance_map;	//variable for the distance-transformed map, type: CV_32FC1
    cv::distanceTransform(map_eroded, distance_map, CV_DIST_L2, 5);
    cv::convertScaleAbs(distance_map, distance_map);	// conversion to 8 bit image

    if(sequence_param.start_pos.size()==3)
    {
        int u= int((sequence_param.start_pos[0]-goal_seq.map_origin.position.x)/goal_seq.map_resolution);
        int v= int((sequence_param.start_pos[1]-goal_seq.map_origin.position.y)/goal_seq.map_resolution);
        if (map_eroded.at<uchar>(v,u) != 0 && distance_map.at<uchar>(v,u) > 15)
        {
            goal_seq.robot_start_coordinate.position.x = sequence_param.start_pos[0];
            goal_seq.robot_start_coordinate.position.y = sequence_param.start_pos[1];
            std::cout<<"use real start_pos" <<std::endl;
            flag=true;
        } 
        else
            std::cout<<"try find a start_pos" <<std::endl;   
    }
    if(!flag)
    {
        bool robot_start_coordinate_set = false;
        for (int v=0; v<map_eroded.rows && robot_start_coordinate_set==false; ++v)
            for (int u=0; u<map_eroded.cols && robot_start_coordinate_set==false; ++u)
                if (map_eroded.at<uchar>(v,u) != 0 && distance_map.at<uchar>(v,u) > 15)
                {
                    goal_seq.robot_start_coordinate.position.x = u*goal_seq.map_resolution + goal_seq.map_origin.position.x;
                    goal_seq.robot_start_coordinate.position.y = v*goal_seq.map_resolution + goal_seq.map_origin.position.y;
                    robot_start_coordinate_set = true;
                }
        std::cout<<"find a start_pos" <<std::endl;
    }
    result_seq= nullptr;
    result_seq= room_sequence_client->sendGoal(goal_seq);
    if(result_seq== nullptr)
    {
        printf("Timeout on room segmentation.");
        return false;
    }
    printf("sequence planning successfully!");

    if(sequence_param.display)
    {
        cv_bridge::CvImagePtr cv_ptr_obj;
        cv_ptr_obj = cv_bridge::toCvCopy(result_seq->sequence_map, sensor_msgs::image_encodings::MONO8);
        cv::imshow("sequence client: ", cv_ptr_obj->image);
        cv::waitKey(0);
    }
    return true;
}

bool IpaRoomPlanning::getRobotPose(std::shared_ptr<tf2_ros::Buffer> tf_,std::vector<double>& robot_pos)
{
    if(tf_==nullptr)
        return false;
    geometry_msgs::msg::PoseStamped global_pose,robot_pose_msg ;
    global_pose.header.frame_id = "map";
    robot_pose_msg.header.frame_id = "base_link";
    tf2::Quaternion quat;
    double roll, pitch, yaw;
    double transform_timeout= 5;
    try 
    {
        tf2::toMsg(tf2::Transform::getIdentity(), global_pose.pose);
        tf2::toMsg(tf2::Transform::getIdentity(), robot_pose_msg.pose);
        robot_pose_msg.header.stamp = rclcpp::Time();
        global_pose= tf_->transform(robot_pose_msg,"map",tf2::durationFromSec(transform_timeout));
        tf2::Quaternion quat = tf2::Quaternion(global_pose.pose.orientation.x, global_pose.pose.orientation.y,
                        global_pose.pose.orientation.z, global_pose.pose.orientation.w);
        tf2::Matrix3x3(quat).getRPY(roll, pitch, yaw);
        robot_pos.clear();
        robot_pos.push_back(global_pose.pose.position.x);
        robot_pos.push_back(global_pose.pose.position.y);
        robot_pos.push_back(yaw);
        return true;
    }
    catch (tf2::LookupException & ex) 
    {
        RCLCPP_ERROR(
        node_handle_->get_logger(),
        "No Transform available Error looking up robot pose: %s\n", ex.what());
        return false;
    } catch (tf2::ConnectivityException & ex) {
        RCLCPP_ERROR(
        node_handle_->get_logger(),
        "Connectivity Error looking up robot pose: %s\n", ex.what());
        return false;
    } catch (tf2::ExtrapolationException & ex) {
        RCLCPP_ERROR(
        node_handle_->get_logger(),
        "Extrapolation Error looking up robot pose: %s\n", ex.what());
        return false;
    } catch (tf2::TimeoutException & ex) {
        RCLCPP_ERROR(
        node_handle_->get_logger(),
        "Transform timeout with tolerance: %.4f", transform_timeout);
        return false;
    } 
    catch (tf2::TransformException & ex) 
    {
        RCLCPP_ERROR(
        node_handle_->get_logger(), "Failed to transform from map to base_link");
        return false;
    }

}


bool IpaRoomPlanning::getRoomSegmentation(IpaRoomPlanningParam & segmentation_param)
{
    std::cout<< "run getRoomSegmentation" << std::endl;
    while(rclcpp::ok() && !room_segmentation_client->wait_server(5)) 
    {
        std::cout<<"room_segmentation_client Action server not available after waiting, wait 5s..."<< std::endl;
    }
    printf("room_segmentation_client, sending goal.\n");
    ipa_building_msgs::action::MapSegmentation::Goal goal_seg;
    goal_seg.data_source=goal_seg.IMAGE_DATA;
    goal_seg.input_map = input_image_msg;
    goal_seg.map_origin.position.x =  segmentation_param.map_origin[0];
    goal_seg.map_origin.position.y =  segmentation_param.map_origin[1];
    goal_seg.map_origin.position.z =  segmentation_param.map_origin[2];
    goal_seg.map_resolution =         segmentation_param.map_resolution;
    goal_seg.return_format_in_meter = segmentation_param.return_format_in_meter;
    goal_seg.return_format_in_pixel = segmentation_param.return_format_in_pixel;
    result_seg= nullptr;
    result_seg= room_segmentation_client->sendGoal(goal_seg);
    if(result_seg== nullptr)
    {
        printf("Timeout on room segmentation.\n");
        return false;
    }
    cv_bridge::CvImagePtr cv_ptr_obj;
    cv_ptr_obj = cv_bridge::toCvCopy(result_seg->segmented_map, sensor_msgs::image_encodings::TYPE_32SC1);
    colour_segmented_map = cv_ptr_obj->image.clone();
    colour_segmented_map.convertTo(colour_segmented_map, CV_8U);
    exploration_map = colour_segmented_map.clone();
    cv::cvtColor(colour_segmented_map, colour_segmented_map, CV_GRAY2BGR);
    for(size_t i = 1; i <= result_seg->room_information_in_pixel.size(); ++i)
    {
        int blue = (rand() % 250) + 1;
        int green = (rand() % 250) + 1;
        int red = (rand() % 250) + 1;
        for(size_t u = 0; u < cv_ptr_obj->image.rows; ++u)
        {
            for(size_t v = 0; v < cv_ptr_obj->image.cols; ++v)
            {
                if(cv_ptr_obj->image.at<int>(u,v) == i)
                {
                    colour_segmented_map.at<cv::Vec3b>(u,v)[0] = blue;
                    colour_segmented_map.at<cv::Vec3b>(u,v)[1] = green;
                    colour_segmented_map.at<cv::Vec3b>(u,v)[2] = red;
                }
            }
        }
    }

    printf("room segmentation successfully!");

    if(segmentation_param.display)
    {
        cv::Mat display=colour_segmented_map.clone();
        for(size_t i = 0; i < result_seg->room_information_in_pixel.size(); ++i)
        {
            cv::Point current_center (result_seg->room_information_in_pixel[i].room_center.x, result_seg->room_information_in_pixel[i].room_center.y);
            cv::circle(display, current_center, 2, CV_RGB(0,0,255), cv::FILLED);
        }
        cv::imshow("segmentation client: ", display);
        cv::waitKey(0);
    }
    return true;
}

